package main

import (
	"context"
	"fmt"
	"github.com/jinzhu/gorm"
	_ "github.com/jinzhu/gorm/dialects/mysql"
	"github.com/silenceper/log"
	"github.com/youtube/vitess/go/pools"
	"time"
)

var MysqlPool *pools.ResourcePool

type MysqlResource struct {
	DB *gorm.DB
}

func (r MysqlResource) Close() {
	r.DB.Close()
}

func NewMysqlConnectionPool() *pools.ResourcePool {
	return pools.NewResourcePool(createResource, 1, 1, time.Millisecond*3600)
}

func createResource() (pools.Resource, error) {
	db, err := connectToMysql()
	return MysqlResource{DB: db}, err
}

func connectToMysql() (db *gorm.DB, err error) {
	user := "root"
	pwd := "root@dev"
	host := "127.0.0.1"
	port := 3306
	dbName := "mysql"
	address := fmt.Sprintf("%s:%s@tcp(%s:%d)/%s?charset=utf8mb4,utf8&parseTime=True&loc=Local", user,
		pwd, host, port, dbName)
	db, err = gorm.Open("mysql", address)
	if err != nil {
		log.Errorf("connect mysql error :%s", err)
	}
	return
}

func getDB() (*gorm.DB, error) {
	if MysqlPool == nil || MysqlPool.IsClosed() {
		// 重新打开资源池
		MysqlPool = NewMysqlConnectionPool()
	}
	r, err := MysqlPool.Get(context.Background())
	if err != nil {
		log.Errorf("connect error:%s", err)
	}
	mr, ok := r.(MysqlResource)
	if !ok || mr.DB == nil {
		db, err := connectToMysql() // 重连
		mr = MysqlResource{DB: db}
		if err != nil {
			log.Errorf("error:%s", err)
			MysqlPool.Put(MysqlResource{DB: &gorm.DB{}}) // 放回去，下次重连
		}
	}
	mr.DB.LogMode(true)
	return mr.DB, err
}

func init() {
	MysqlPool = NewMysqlConnectionPool()
}

func main() {
	usePoolTest()
	useNoPoolTest()
}

func mysqlPoolTest() (count int, err error) {
	db, err := getDB()
	defer MysqlPool.Put(MysqlResource{DB: db})
	if err != nil {
		log.Error(err.Error())
		return
	}
	err = db.Table("user").Count(&count).Error
	if err != nil {
		log.Error(err.Error())
	}
	log.Infof("%d", count)
	return
}

func mysqlTest() (count int, err error) {
	db, err := connectToMysql()
	defer db.Close()
	if err != nil {
		log.Error(err.Error())
		return
	}
	err = db.Table("user").Count(&count).Error
	if err != nil {
		log.Error(err.Error())
	}
	log.Infof("%d", count)
	return
}

func useNoPoolTest() {
	startTime := time.Now()
	defer func() {
		log.Infof("no use pool:%v", float64(time.Now().Sub(startTime))/float64(time.Second))
	}()
	for i := 1; i < 30; i++ {
		mysqlTest()
	}
}

func usePoolTest() {
	startTime := time.Now()
	defer func() {
		log.Infof("use pool:%v", float64(time.Now().Sub(startTime))/float64(time.Second))
	}()
	for i := 1; i < 30; i++ {
		mysqlPoolTest()
	}
}
